import paramiko
import yaml
import os
from contextlib import contextmanager

class SSHManager:
    def __init__(self):
        self.config = self._load_config()
        self.ssh = None
        
    def _load_config(self):
        """加载SSH配置"""
        config_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'config', 'config.yaml')
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                return yaml.safe_load(f)
        except Exception as e:
            print(f"加载配置文件失败: {str(e)}")
            return {
                "servers": {
                    "server_a": {
                        "host": "localhost",
                        "port": 22,
                        "username": "user",
                        "password": "password"
                    }
                },
                "paths": {
                    "source_path": "/path/to/source",
                    "destination_path": "/path/to/destination"
                }
            }
    
    def _create_ssh_client(self, server_config):
        """创建SSH客户端连接"""
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            ssh.connect(
                hostname=server_config['host'],
                port=server_config['port'],
                username=server_config['username'],
                password=server_config['password']
            )
            return ssh
        except Exception as e:
            raise Exception(f"SSH连接失败: {str(e)}")
    
    @contextmanager
    def create_tunnel(self):
        """创建SSH隧道"""
        ssh_d_to_c = None
        ssh_c_to_b = None
        ssh_b_to_a = None
        
        try:
            # D -> C
            print("正在连接服务器C...")
            ssh_d_to_c = self._create_ssh_client(self.config['servers']['server_c'])
            transport_d_to_c = ssh_d_to_c.get_transport()
            
            # C -> B
            print("正在连接服务器B...")
            local_port_c = 10022
            channel_c = transport_d_to_c.open_channel(
                "direct-tcpip",
                (self.config['servers']['server_b']['host'], self.config['servers']['server_b']['port']),
                ('localhost', local_port_c)
            )
            ssh_c_to_b = paramiko.SSHClient()
            ssh_c_to_b.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            ssh_c_to_b.connect(
                'localhost',
                port=local_port_c,
                username=self.config['servers']['server_b']['username'],
                password=self.config['servers']['server_b']['password'],
                sock=channel_c
            )
            
            # B -> A
            print("正在连接服务器A...")
            local_port_b = 10023
            channel_b = ssh_c_to_b.get_transport().open_channel(
                "direct-tcpip",
                (self.config['servers']['server_a']['host'], self.config['servers']['server_a']['port']),
                ('localhost', local_port_b)
            )
            ssh_b_to_a = paramiko.SSHClient()
            ssh_b_to_a.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            ssh_b_to_a.connect(
                'localhost',
                port=local_port_b,
                username=self.config['servers']['server_a']['username'],
                password=self.config['servers']['server_a']['password'],
                sock=channel_b
            )
            
            print("所有服务器连接成功！")
            yield ssh_b_to_a
            
        except Exception as e:
            print(f"连接错误: {str(e)}")
            raise
        finally:
            if ssh_b_to_a:
                ssh_b_to_a.close()
            if ssh_c_to_b:
                ssh_c_to_b.close()
            if ssh_d_to_c:
                ssh_d_to_c.close()
    
    def execute_command(self, command):
        """执行SSH命令并返回结果"""
        try:
            with self.create_tunnel() as ssh:
                stdin, stdout, stderr = ssh.exec_command(command)
                output = stdout.read().decode('utf-8')
                error = stderr.read().decode('utf-8')
                exit_status = stdout.channel.recv_exit_status()
                
                return {
                    'status': exit_status,
                    'output': output,
                    'error': error
                }
        except Exception as e:
            print(f"执行命令失败: {str(e)}")
            return {
                'status': -1,
                'output': '',
                'error': str(e)
            }
